"""
数据预处理
"""
import numpy as np
import pandas as pd


def series_2_train_data(data, n_out=1, dropnan=True):
    """
    将时序数据转为训练集：特征部分 & 预测部分
    :param data:
    :param n_out: 预测的steps
    :param dropnan: 是否删除nan数据
    :return:
    """
    features = 1 if isinstance(data, list) else data.shape[1]
    # 特征列
    names = ['var(t-{})'.format(i) for i in
             range(features, 0, -1)]

    # 预测列
    forecast_data = data[:, 0].tolist()
    for i in range(n_out):
        # 列名合并
        names += ['var(t)'] if i == 0 else ['var(t+{})'.format(i)]

        forcast_columns = forecast_data[i + 1:]
        [forcast_columns.append(np.nan) for _ in range(i + 1)]
        # 数据合并
        data = np.column_stack((data, forcast_columns))

    df_train = pd.DataFrame(data[: -1 * n_out], columns=names)
    if dropnan:
        df_train = df_train.dropna()

    return df_train
